import qualified Data.Map as Map 
import Data.List (isPrefixOf, findIndex)
import Data.Char (isDigit)

data MaskKind = Keep | Set | Floating deriving (Show)
type Mask = [MaskKind]

type Address = Integer
type Value = Integer

data Instruction = SetMask Mask 
                 | WriteMem Address Value deriving (Show)

type Memory = Map.Map Address Value
data Computer = Computer Mask Memory deriving (Show)

parseMaskKind :: Char -> MaskKind
parseMaskKind c = case c of
    '0' -> Keep
    '1' -> Set
    'X' -> Floating

-- mask = X11001110001101XX01111X1001X01101111
-- mem[32163] = 23587
parseInstruction :: String -> Instruction
parseInstruction string 
    | isPrefixOf "mask" string = SetMask mask
    | isPrefixOf "mem" string  = WriteMem address value
    | otherwise = error "Unkonw instruction!"
        where
            mask = map parseMaskKind . drop 7 $ string
            address = read . takeWhile isDigit . drop 4 $ string
            value = read . drop 2 . dropWhile (/='=') $ string

parseInput :: String -> [Instruction]
parseInput = map parseInstruction . lines

binaryToAddress :: String -> Address
binaryToAddress = foldl foldingFunc 0
    where
        foldingFunc :: Address -> Char -> Address
        foldingFunc address ch = case ch of
                '0' -> address * 2
                '1' -> address * 2 + 1

addressToBinary :: Address -> String
addressToBinary address = (replicate len '0') ++ string
    where
        string = reverse . addressToBinaryRev $ address
        len = 36 - length string
    
        addressToBinaryRev :: Address -> String
        addressToBinaryRev 1 = "1"
        addressToBinaryRev 0 = "0"
        addressToBinaryRev address = (if reminder == 1 then '1' else '0') : addressToBinaryRev half
            where
                reminder = address `mod` 2
                half = address `div` 2

withMask :: MaskKind -> Char -> Char
withMask m c = case m of
    Keep -> c
    Set  -> '1'
    Floating -> 'X'

replaceIndex :: String -> Int -> Char -> String
replaceIndex (s:ss) 0 ch = ch:ss
replaceIndex (s:ss) idx ch = s : (replaceIndex ss (idx-1) ch)

expandFloating :: String -> [String]
expandFloating string = case floatingIndex of
    Nothing -> [string]
    Just index -> expandFloating replace1 ++ expandFloating replace2
        where
            replace1 = replaceIndex string index '1'
            replace2 = replaceIndex string index '0'
    where 
        floatingIndex = findIndex (=='X') string

maskAddress :: Mask -> Address -> [Address]
maskAddress mask address = map binaryToAddress binStrs'
    where
        binStr = addressToBinary address
        binStrs' = expandFloating . zipWith withMask mask $ binStr

executeInstruction :: Computer -> Instruction -> Computer
executeInstruction (Computer _ memory) (SetMask mask) = Computer mask memory
executeInstruction (Computer mask memory) (WriteMem address value) = Computer mask memory'
    where
        addresses = maskAddress mask address
        memory' = foldl (\mem addr -> Map.insert addr value mem) memory addresses

execute :: [Instruction] -> Computer 
execute = foldl executeInstruction (Computer (repeat Keep) Map.empty)

memorySum :: Computer -> Value
memorySum (Computer _ memory) = sum . map snd . Map.toList $ memory

main :: IO ()
main = do
    input <- readFile "input.txt"
    print . memorySum . execute . parseInput $ input